import torch
import torch.nn as nn
import torch.nn.functional as F

# 假设有3个类别
from torch import optim

num_classes = 3

# 假设有一个输入特征张量，batch_size = 5
inputs = torch.randn(5, requires_grad=True)

# 定义一个全连接层
linear = nn.Linear(in_features=inputs.size(1), out_features=num_classes)
optimizer = optim.SGD(linear.parameters(), lr=0.01)  # 学习率设置为0.01
# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 进行线性变换
logits = linear(inputs)
# 假设目标标签是一个包含类别索引的张量
targets = torch.tensor([0, 2, 1, 0, 1])
# 计算损失
loss = criterion(F.log_softmax(logits, dim=1), targets)

# 反向传播
loss.backward()
# 清除梯度，以便进行下一次迭代
optimizer.zero_grad()
# 更新梯度
optimizer.step()

